#include <stdio.h>
#include <iostream>
#include <stdlib.h>
#include <string.h>
#include <unistd.h>
#include <sys/types.h>
#include <sys/socket.h>
#include <netinet/in.h>
#include <sys/ioctl.h>
#include <errno.h>
#include "Util.h"
#include "Polling.h"
#include "Keyboard.h"
#include <time.h>

/*---------------------------------------------------------------
    DEFINES
---------------------------------------------------------------*/
#define MAX_APPS 4 
#define MAX_SOCKETS 32 
#define MAX_PACKET 2048

#define STAT_DEAD 0
#define STAT_RUN 1

#define MSG_PEER_MSG 1
#define MSG_PEER_DIED 2
#define MSG_SERVER_DIED 3

using namespace std;

/*---------------------------------------------------------------
    GLOBAL VARIABLES 
---------------------------------------------------------------*/

int sockFD, connectFD;
int sockets[MAX_APPS][MAX_SOCKETS];
int socketStatus[MAX_APPS][MAX_SOCKETS];
int initSockets[MAX_SOCKETS];
int appStatus[MAX_APPS];
int appSocketCount[MAX_APPS];
int appServerID[MAX_APPS];
int totalSockets;

struct sockaddr_in serverAddress, clientAddress;

char logFile[256];

/*---------------------------------------------------------------
    PREDECLARATIONS
---------------------------------------------------------------*/
int sendAll(int appID, void *buff, int size, int except);

/*-------------------------------------------
    log(...)
-------------------------------------------*/
void log(int appID, int socketID, const char *message)
{
    time_t t;
    time(&t); 

    char *time = ctime(&t);
    time[strlen(time)-1] = 0;
       
    cout << time << " | " << appID << ":" << socketID << " | " << message << endl;
}

/*-------------------------------------------
    socketDeathHandler(...)
-------------------------------------------*/
int socketDeathHandler (int appID, int socketDead)
{
    if (socketStatus[appID][socketDead] == STAT_DEAD)
        return 0;

    totalSockets--;

    sockets[appID][socketDead] = STAT_DEAD;
    socketStatus[appID][socketDead] = STAT_DEAD;
    appSocketCount[appID]--;

    char msg[2];

    if (appServerID[appID] != socketDead)
    {
        // Communicate that a non-server peer 'died'
        msg[0] = MSG_PEER_DIED;
        msg[1] = socketDead + 1;
        log(appID, socketDead, "Connection Lost. Echoing Death to Peers.");
    }
    else
    {
        log(appID, socketDead, "Server Lost. Echoing Death and Finding new server.");

        // Communicate that the server died.
        msg[0] = MSG_SERVER_DIED;
        msg[1] = 0;

        // Find who is going to be the next server.
        for (int i = 0; i < MAX_SOCKETS; i++)
        {
            if (socketStatus[appID][i] != STAT_DEAD)
            {
                log(appID, i, "Found new server.");

                msg[1] = (i + 1);
                appServerID[appID] = i;
                break;
            }
        }
    }

    // Ping the message to all servers of the peer/server death.
    sendAll(appID, msg, 2, 0);

    return 1;
}

/*-------------------------------------------
    sendAll(...)
-------------------------------------------*/
int sendAll(int appID, void *buff, int size, int except)
{
    for (int i = 0; i < MAX_SOCKETS; i++)
    {
        // We don't echo data to the client that sent it !
        if (i == except)
           continue;

        if (socketStatus[appID][i] != STAT_DEAD)
        {
             if (!write(sockets[appID][i], buff, size))
             {
                 return socketDeathHandler(appID, i);
             }
        }
    }

    return 1;
}


/*-------------------------------------------
    acceptConnection(...)
-------------------------------------------*/
void acceptConnection()
{
    socklen_t len = sizeof(clientAddress);

    if (pollRead(sockFD))
    {
        connectFD = accept(sockFD, (struct sockaddr *) &clientAddress, &len);

        if (connectFD < 0)
        {
            log(-1,-1, "!! Error occurred accepting connection. ");
            return;
        }

        totalSockets++;

        for (int i = 0; i < MAX_SOCKETS; i++)
        {
            if (initSockets[i] == STAT_DEAD)
            {
                log(-1, i, "New Socket Connection. Waiting for applicationID."); 

                initSockets[i] = connectFD;
                break;
            }
        }
    }
}

/*-------------------------------------------
    readApplicationID(...)
-------------------------------------------*/
void readApplicationID(int initIX)
{
    char c_appID;
    int n = read(initSockets[initIX], &c_appID, 1);

    if (n <= 0)
    {
        log(-1,-1,"New connection closed before appID received.");

        // This socket died during initialization
        initSockets[initIX] = STAT_DEAD;
        totalSockets--;
        return;
    }

    int appID = c_appID;

    log(appID, -1, "appID received.");

    for (int s = 0; s < MAX_SOCKETS; s++)
    {
        if (socketStatus[appID][s] == STAT_DEAD)
        {
            log(appID, s, "clientID found! New client connected and now running.");

            sockets[appID][s] = initSockets[initIX];
            socketStatus[appID][s] = STAT_RUN;
            appStatus[appID] = STAT_RUN;
           
            char clientID = s + 1; 
            write(sockets[appID][s], &clientID, sizeof(char));

            // Is this the first client? They are now server.
            // We communicate this with a 0
            if (appSocketCount[appID] <= 0)
            {
                char isFirst = s + 1;
                appServerID[appID] = s;
                write(sockets[appID][s], &isFirst, 1); 
                appSocketCount[appID] = 1;
            }
            else
            {
                char isFirst = appServerID[appID] + 1;
                write(sockets[appID][s], &isFirst, 1); 
                appSocketCount[appID]++;
            }

            break;
        }
    }

    initSockets[initIX] = STAT_DEAD;
}

/*-------------------------------------------
    readSocket(...)
-------------------------------------------*/
int readSocket(int a, int s)
{
    int nBytes;
    ioctl(sockets[a][s], FIONREAD, (char *)&nBytes);

    char bytes[nBytes];
    bzero(bytes, nBytes);

    int n = read(sockets[a][s], bytes, nBytes);

    if (n <= 0)
    {
        return socketDeathHandler(a, s);
    }

    char data[4 + (n-1)];
    data[0] = 1; // This communicates that this is a data sequence
    data[1] = (char)(s + 1); // This is the from peer ID 
    short *count = (short *)&data[2];

    *count = (short)(n - 1); // This communicates the number of bytes received by the client

    for (int b = 1; b < n; b++)
    {
        data[3 + b] = bytes[b];
    }

    // Do we send to everyone?
    if (bytes[0] == 0)
    {
        // Send to everyone except s
        return sendAll(a, data, 4 + (n-1), s);
    }

    // clientIDs are sent 1-based to the peers. We have to convert
    // back into 0-based for our array.
    int j = bytes[0] - 1;

    if (socketStatus[a][j] != STAT_DEAD)
    {
        if (!write(sockets[a][j], data, 4 + (n-1)))
        {
            return socketDeathHandler(a, j);
        }
    }

    // Echo success
    return 1;
}

/*-------------------------------------------
    readClients(...)
-------------------------------------------*/
void readClients()
{
    for (int is = 0; is < MAX_SOCKETS; is++)
    {
        if (initSockets[is] != STAT_DEAD && pollRead(initSockets[is]))
        {
            readApplicationID(is);
        }
    }

    for (int a = 0; a < MAX_APPS; a++)
    {
        if (appStatus[a] == STAT_DEAD)
            continue;

        for (int s = 0; s < MAX_SOCKETS; s++)
        {
            if (socketStatus[a][s] != STAT_DEAD && pollRead(sockets[a][s]))
            {
                switch (socketStatus[a][s])
                {
                    case STAT_RUN:
                        readSocket(a, s);
                    break;
                }
            }
        }
    }
}

/*-------------------------------------------
    sleepIfIdle(...)
-------------------------------------------*/
void sleepIfIdle()
{
    if (totalSockets <= 0)
    {
        totalSockets = 0;
        sleep(1);
    }
}

/*-------------------------------------------
    startServer(...)
-------------------------------------------*/
void startServer(ushort port)
{
    sockFD = socket(AF_INET, SOCK_STREAM, 0);

    if (sockFD < 0)
    {
        log(-1,-1, "FATAL Error opening socket.");
        exit(-1);
	}

    serverAddress.sin_family = AF_INET;
    serverAddress.sin_addr.s_addr = htonl(INADDR_ANY);
    serverAddress.sin_port = port;

    if (bind(sockFD, (struct sockaddr *) &serverAddress, sizeof(struct sockaddr_in)) < 0)
    {
        char buff[256];
        sprintf(buff,"FATAL Error binding to socket for port %d: %s", port, strerror(errno));

        log(-1,-1, buff);
        exit(-2);
    }

    listen(sockFD, 5);

    cout << "Serving starting..." << endl;
    cout << "Sending to log file /var/log/echosurf ..." << endl;

    startLog(logFile);

    log(-1,-1, "------------------------------- STARTED --------------------------------");

    while (true)
    {
        //usleep(500);

        sleepIfIdle();

        acceptConnection();

        readClients();
    }
}

/*-------------------------------------------
    main(...)
-------------------------------------------*/
int main(int argc, char **argv)
{
    for (int is = 0; is < MAX_SOCKETS; is++)
    {
        initSockets[is] = STAT_DEAD;
    }

    for (int a = 0; a < MAX_APPS; a++)
    {
        appStatus[a] = STAT_DEAD;
        appServerID[a] = -1;
        appSocketCount[a] = 0;

        for (int s = 0; s < MAX_SOCKETS; s++)
        {
            sockets[a][s] = STAT_DEAD;
            socketStatus[a][s] = STAT_DEAD;
        }
    }
   
    // Start the daemon! 
    daemonize();

    unsigned short port = 1002;

    if (argc == 2)
    {
        port = atoi(argv[1]);
    }

    char buff[256];
    sprintf(buff,"Starting EchoServe on port: %d", port);
 
    log(-1,-1, buff);
    
    // Convert to little endian 
    endianSwap(port);

    bzero(logFile, 256);
    strcpy(logFile,"/var/log/echosurf");  

    if (argc == 3)
    {
        strcpy(logFile, argv[2]);
    }

    startServer(port);
}
